#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy, h5py
import time
import csv
from pyscf_ccso import py2cc,  mo_sym, dfoutcore, ri_blocks, canonical
from pyscf import gto, lib, scf, symm, lib, ao2mo, df
#====================================================================================================  
#====================================================================================================
def rhf(myhf, mol, nf=0, start_nbo_pyscf=0, thresh=0.00, start_ncc_cfour=0, dropmo_cf=0, get_eri=1,
        auxbas=None):
# 读取CCSOINP中的相关信息 
   iop_nbo_state, iop_nbo_act, iop_ri, iop_soc, _, eri_level = py2cc.readinp() 

# 得到点群名称，nirrep，irrep name and idx 
   group_corr, nirrep, irrep_sym = py2cc.symminfo(mol.groupname)

# symrx, symry and symrz for scalar relativistic 
   symrx=1; symry=1; symrz=1
# orbit information
   pop, vrt, isymmo = py2cc.symorb_sort(mol, nirrep, myhf.mo_occ, myhf.mo_coeff,nf)

# get dirprd 
   dirprd = py2cc.get_dirprd()

# 由密度矩阵得到。自然轨道 
   if(start_nbo_pyscf==1):
      vrt, pop, isymmo = py2cc.start_nat_orb(mol, myhf, nirrep, thresh, nf)
# save symm data
   nmo = sum(pop) + sum(vrt) 
   mo_sym.save_symm(dirprd, nirrep, nmo, vrt, pop, isymmo, symrx, symry, symrz)
# save eri
   py2cc.generate_eri(mol, myhf, pop, vrt, nf, nirrep, iop_ri, eri_level)

# 存储分子轨道能量,占据轨道与空轨道数目等。
   py2cc.save_cc_pyscf_int(sum(pop), sum(pop), sum(vrt), sum(vrt), nmo, myhf.e_tot, 
                             group_corr, nf, nirrep, irrep_sym, 
                             myhf.mo_energy[nf:], myhf.mo_energy[nf:])
#r
   return nirrep, dirprd, vrt, pop
#================================================================================================================= 
def uhf(myhf, mol, nf=0, start_nbo_pyscf=0, thresh=0.00, start_ncc_cfour=0, dropmo_cf=0, get_eri=1, auxbas=None): 
# 读取CCSOINP中的相关信息 
   iop_nbo_state, iop_nbo_act, iop_ri, iop_soc, _, eri_level = py2cc.readinp()
# 得到点群名称，nirrep，irrep name and idx 
   group_corr, nirrep, irrep_sym = py2cc.symminfo(mol.groupname)
# symrx, symry and symrz for scalar relativistic
   symrx=1; symry=1; symrz=1 
# orbit information
   mo_occ=myhf.mo_occ[0]; mo_coeff=myhf.mo_coeff[0] 
#  print('shape', mo_occ.shape, mo_coeff.shape)
#  print('mo_occ', myhf.mo_occ)
#  print('mo_coeff a', myhf.mo_coeff[0])
#  print('mo_coeff b', myhf.mo_coeff[1])
   pop, vrt, isymmo = py2cc.symorb_sort_u(mol, nirrep, myhf.mo_occ, myhf.mo_coeff,nf)
# get dirprd 
   dirprd = py2cc.get_dirprd()
   nmo = sum(pop[:,0]) + sum(vrt[:,0]) 
# save symm data
   mo_sym.save_symm_uhf(dirprd, nirrep, nmo, vrt, pop, isymmo, symrx, symry, symrz)

# 存储分子轨道能量,占据轨道与空轨道数目等。
   py2cc.save_cc_pyscf_int(sum(pop[:,0]), sum(pop[:,1]), sum(vrt[:,0]), sum(vrt[:,1]),
                           nmo, myhf.e_tot, group_corr, nf, nirrep, irrep_sym, 
                           myhf.mo_energy[0][nf:], myhf.mo_energy[1][nf:])
# save eri
   py2cc.generate_eri_uhf(mol, myhf, pop, vrt, nf, nirrep, iop_ri, eri_level)


def rohf(myhf, mol, nf=0, start_nbo_pyscf=0, thresh=0.00, start_ncc_cfour=0, dropmo_cf=0, get_eri=1, auxbas=None): 
# 读取CCSOINP中的相关信息 
   iop_nbo_state, iop_nbo_act, iop_ri, iop_soc, _, eri_level = py2cc.readinp()
# 得到点群名称，nirrep，irrep name and idx 
   group_corr, nirrep, irrep_sym = py2cc.symminfo(mol.groupname)
# symrx, symry and symrz for scalar relativistic
   symrx=1; symry=1; symrz=1 
# orbit information
   mo_occ=myhf.mo_occ; mo_coeff=myhf.mo_coeff 
   print('shape', mo_occ.shape, mo_coeff.shape)
   print('mo_occ', myhf.mo_occ)
   print('mo_coeff a', myhf.mo_coeff.shape)
   print('mo_coeff b', myhf.mo_coeff.shape)
   print('vrt and pop')
   myhf, mo_coeff_new, mo_energy_new = canonical.canonical_rohf(myhf, mol, mo_coeff, mo_occ, nf, savefvo=True)
#   newmo_coeffa, newmo_coeffb, newmo_ea, newmo_eb = canonical.canonical_rohf(myhf, mol, mo_coeff, mo_occ, savefvo=True)
   pop, vrt, isymmo = py2cc.symorb_sort_ro(mol, nirrep, myhf.mo_occ, myhf.mo_coeff.ro,nf)
# get dirprd 
#  print('vrt', vrt)
#  print('pop', pop)
   dirprd = py2cc.get_dirprd()
   nmo = sum(pop[:,0]) + sum(vrt[:,0]) 
# save symm data
   mo_sym.save_symm_uhf(dirprd, nirrep, nmo, vrt, pop, isymmo, symrx, symry, symrz)
#   print('mo_energy', myhf.mo_energy)
# 存储分子轨道能量,占据轨道与空轨道数目等。
   py2cc.save_cc_pyscf_int(sum(pop[:,0]), sum(pop[:,1]), sum(vrt[:,0]), sum(vrt[:,1]),
                           nmo, myhf.e_tot, group_corr, nf, nirrep, irrep_sym, 
                           myhf.mo_energy.ro[0,nf:], myhf.mo_energy.ro[1,nf:])
# save eri
   py2cc.generate_eri_rohf(mol, myhf, pop, vrt, nf, nirrep, iop_ri, eri_level)
